from app.db.session import get_db
from sqlalchemy.orm import Session
from app.core import security
from sqlalchemy import create_engine, text

from app.models.ai import AIConfig
from app.models.base import BaseEnumYesNo


class DatabaseInitData:

    def __init__(self, db_url: str):
        self.db_url = db_url
        # print(f"alembic_cfg: ")

    def set_auto_increment(self,db:Session):
        # 设置主键初始值
        db.execute(text("ALTER TABLE ai_agents AUTO_INCREMENT = 5000"))
        db.execute(text("ALTER TABLE ai_api AUTO_INCREMENT = 5000"))
        db.execute(text("ALTER TABLE ai_model AUTO_INCREMENT = 5000"))

    def startup(self):
        #初始化用户
        db: Session = next(get_db())
        if not self.is_init(db):
            db.close()
            return
        self.set_auto_increment(db)
        self.init_ai_config(db)
        self.init_user(db)

        db.commit()
        db.close()

    def is_init(self,db:Session):
        _count=db.query(AIConfig).count()
        return _count==0

    def init_user(self,db:Session):
        from app.models.user import User
        from app.schemas.user import UserCreate
        user_in= UserCreate(username='admin',email="xxx@xx.com",password="123456",nickname="AI程序员")
        user = db.query(User).filter(User.username == user_in.username).first()
        if user is None:
            user = User(
                email=user_in.email,
                username=user_in.username,
                hashed_password=security.get_password_hash(user_in.password),
                nickname=user_in.nickname
            )
            db.add(user)
            db.refresh(user)

    def init_ai_config(self,db:Session):
        _count=db.query(AIConfig).count()
        if _count==0:
            db.add_all(
                [
                    AIConfig(
                        config_key="ai.defaultModel",
                        config_name="默认模型",
                        config_value="",
                        description="默认模型",
                        remark="创建新助手时使用的模型，如果助手未设置模型，则使用此模型",
                        is_system=BaseEnumYesNo.YES.value.value,
                        user_id=1,
                    ),
                    AIConfig(
                        config_key="ai.topicNamingModel",
                        config_name="话题命名模型",
                        config_value="",
                        remark="自动命名新话题时使用的模型",
                        is_system=BaseEnumYesNo.YES.value.value,
                        user_id=1,
                    ),
                    AIConfig(
                        config_key="ai.translateModel",
                        config_name="翻译模型",
                        config_value="",
                        remark="翻译服务使用的模型",
                        is_system=BaseEnumYesNo.YES.value.value,
                        user_id=1,
                    ),
                ]
            )
